import json
import glob
import os
data_dir = "/data2/qbl/ccta_add"
all_images = sorted(glob.glob(os.path.join(data_dir, "image_100", "*.nii.gz")))
all_labels = sorted(glob.glob(os.path.join(data_dir, "vessel_largest_100", "*.nii.gz")))
for i in range(len(all_images)):
    all_images[i] = os.path.join(*all_images[i].split("/")[-2:])
    all_labels[i] = os.path.join(*all_labels[i].split("/")[-2:])
train_images, val_images, test_images = all_images[:80], all_images[80:90],all_images[90:]
train_labels, val_labels = all_labels[:80], all_labels[80:90]
train_list = []
val_list = []
for i in range(len(train_images)):
    train_list.append({"image":train_images[i],"label":train_labels[i]})
for i in range(len(val_images)):
    val_list.append({"image":val_images[i],"label":val_labels[i]})
content_dict = {
  "description": "ASACA100",
  "labels": {
      "0": "background",
      "1": "artery",
      "2": "coronary"
  },
  "licence": "yt",
  "modality": {
      "0": "CT"
  },
  "name": "ASACA100",
  "numTest": 10,
  "numTraining": 90,
  "reference": "",
  "release": "",
  "tensorImageSize": "3D",
  "test": [filename for filename in test_images],
  "training": train_list,
  "validation": val_list,
}

with open(os.path.join(data_dir, 'ASACA100.json'), 'w') as json_file:
    json.dump(content_dict, json_file)

